
#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <mpi.h>
#include "jacobi.h"
#include "matrixparser.h"



double max_openmp(double *vdif, double *vx, int size)
{
    double max,max2;

    max = max2 = 0;

    omp_set_num_threads(1);
    #pragma omp parallel
    #pragma omp sections
    {
        #pragma omp section
            max = diff(vdif,size);
        #pragma omp section
            max2 = diff(vx,size);
    }

    return max/max2;
}

double diff(double *diff, int size)
{
    int i;
    double max = 0;
    double temp;

    for(i=0;i<size;i++)
    {
        temp = fabs(diff[i]);

        if(temp>max)
            max = temp;
    }

    return max;
}

double *solve_mpi(int argc, char *argv[],double *a, double *b, double *x, int m_order, int row_test,int max_it, double eps)
{
    unsigned int m_order_pow_2 = m_order * m_order;
    int j = 0;
    int i = 0;
    int it;
    int lines_per_task;
    int lines_tail;
    int index;
    int indexAdjust;
    double *recvbuf;
    double *bbuf;
    double diagonal_temp;
    int tam;
    int tail;
    double *xTmp = malloc(m_order * sizeof(double));
    int t;
    int loop = 1;
    int convergence = 1;
    omp_set_num_threads(1);
  
    double old_b;
    double *old_a;

    old_a = malloc(m_order * sizeof(double));
    old_b = b[row_test];
    for(i = 0; i < m_order; i++)
    {
        old_a[i] = a[m_order*row_test + i];
    }

    /*mpi declarations*/
    int rank,source,numtasks,sendcount,recvcount;
    double *dif = malloc(m_order*sizeof(double));
    double *dbuf;    
    MPI_Init(&argc,&argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &numtasks);
    
    source = 0;

   

    /*Calculate sizes*/
    if(numtasks > m_order)
    {
        lines_per_task = m_order;
        tam = m_order;
        lines_tail = 0;
    }else {
        lines_per_task = m_order/numtasks;
        lines_tail = m_order%numtasks;
        tam = lines_per_task * m_order;
    }


    tail = lines_tail * m_order;
    sendcount = tam;
    recvcount = tam;

    recvbuf = malloc(recvcount* sizeof(double));
    bbuf = malloc(lines_per_task * sizeof(double));
    dbuf = malloc(lines_per_task*sizeof(double));


    /*Send lines to tasks to process*/
    MPI_Scatter(a,sendcount,MPI_DOUBLE,recvbuf,recvcount,
             MPI_DOUBLE,source,MPI_COMM_WORLD);
    MPI_Scatter(b,lines_per_task,MPI_DOUBLE,bbuf,lines_per_task,
             MPI_DOUBLE,source,MPI_COMM_WORLD);
    MPI_Scatter(dif,lines_per_task,MPI_DOUBLE,dbuf,lines_per_task,
             MPI_DOUBLE,source,MPI_COMM_WORLD);
    

    /*Process lines divided for each process*/
    for(i = 0; i < lines_per_task; i++)
    {
        index = lines_per_task*rank;
        index = index +((m_order+1)*i);
        diagonal_temp = recvbuf[index];
        recvbuf[index] = 0;
        indexAdjust = i*m_order;

        for(j = 0; j < m_order;j++)
        {
            if(j != index)
                recvbuf[j+indexAdjust] /= diagonal_temp;
        }

        bbuf[i] /= diagonal_temp;
        if(row_convergence(i, recvbuf, m_order, eps) == 0)
            convergence = 0;
    }
   
    /*Process lines from the tail*/
    if(rank == source && tail != 0)
    {
        for(i = 0; i < lines_tail; i++)
        {
           
            index = m_order_pow_2 - (i*(m_order+1)) - 1;
            indexAdjust =  m_order_pow_2 - ((i+1)*m_order);
            diagonal_temp = a[index];
            a[index] = 0;
            int c = m_order - i - 1; 

            for(j = m_order - 1; j >= 0; j--)
            {
                if(j != c)
                    a[j+indexAdjust] /= diagonal_temp;
            }

            b[m_order - i - 1] /= diagonal_temp;

            if(row_convergence(m_order - i - 1, a, m_order, eps) == 0)
                convergence = 0;
        }
    } 


    MPI_Allreduce(&convergence,&loop,1, MPI_INT,MPI_MIN,MPI_COMM_WORLD);
    
    if(rank == source)
    {
        if(loop == 1)
            printf("Convergence test passed\n");
        else
            printf("Convergence teste failed\n");
    }

    MPI_Barrier(MPI_COMM_WORLD);
    omp_set_num_threads(1);

    /*BEGIN Main loop to find result*/
    for(it = 0; it < max_it && loop == 1; it++)
    {
        /*Broadcast last Xk values to all tasks*/
        MPI_Bcast(x,m_order,MPI_DOUBLE,source,MPI_COMM_WORLD);
    
        for(i = 0; i < m_order; i++)
            xTmp[i] = x[i];


        if(rank == source)
        {
            #pragma omp parallel shared(x) private (i,j,indexAdjust)
           {
            #pragma omp sections nowait
            {
                #pragma omp section
                    for(i = 0; i < lines_per_task; i++)
                    {
                        index = (rank*lines_per_task) + i;
                        indexAdjust = i*m_order;
                        dbuf[i] = xTmp[index];
                        x[index] = bbuf[i];

                        for(j = 0; j < m_order;j++)
                        {
                            x[index] -= recvbuf[j+indexAdjust]*xTmp[j];
                        }
                        dbuf[i] -= x[index];
                    }

                #pragma omp section
                    if(tail != 0)
                    {
                         for(i = 0; i < lines_tail; i++)
                         {
                            indexAdjust =  m_order_pow_2 - ((lines_tail - i)*m_order);
                            int c = m_order - lines_tail + i;
                            dif[c]  = xTmp[c];  
                            x[c] = b[c];


                            for(j = 0; j < m_order; j++)
                            {
                                x[c] -= a[j+indexAdjust]*xTmp[j];
                            }   
                            dif[c] -= x[c];
                        }
                    }
            }
        }
    }
        
       else 
        {
            #pragma omp parallel for firstprivate(j) lastprivate(i)
            for(i = 0; i < lines_per_task; i++)
            {
                index = (rank*lines_per_task) + i;
                indexAdjust = i*m_order;
                dbuf[i] = xTmp[index];
                x[index] = bbuf[i];
                for(j = 0; j < m_order;j++)
                {
                    x[index] -= recvbuf[j+indexAdjust]*xTmp[j];  
                }
                dbuf[i] -= x[index];
            }
        }

        /*Gather result from dif and x from each process*/
        MPI_Gather( dbuf, lines_per_task, MPI_DOUBLE, 
            dif,lines_per_task, MPI_DOUBLE, source, MPI_COMM_WORLD);
        MPI_Gather( &x[rank*lines_per_task], lines_per_task, MPI_DOUBLE, 
            x,lines_per_task, MPI_DOUBLE, source, MPI_COMM_WORLD);



        if(rank == source)
        {
            
            if(max_openmp(dif,x,m_order) <= eps)
            {
                printf("Iterations %d\n",it);
                loop = 0;
            }
        }

        MPI_Barrier(MPI_COMM_WORLD);
        MPI_Bcast(&loop,1,MPI_INT,source,MPI_COMM_WORLD);
    }
    /*END Main loop*/

    /*Get the values for A* and B* from other tasks*/
    MPI_Gather( recvbuf, recvcount, MPI_DOUBLE, 
            a,sendcount, MPI_DOUBLE, source, MPI_COMM_WORLD);
    MPI_Gather( bbuf, lines_per_task, MPI_DOUBLE, 
            b,lines_per_task, MPI_DOUBLE, source, MPI_COMM_WORLD);
        


    /*Check result*/
    if(rank == source && it > 0)
    {
        check_result(row_test,old_a,x,old_b, m_order);
    }

    MPI_Finalize();
    return x;
}


double *solve(double *a, double *b, double *x, int m_order, int row_test,int max_it, double eps)
{

    unsigned int m_order_pow_2 = m_order * m_order;
    int temp = 0;
    int cond = 0;
    int j = 0;
    int i = 0;
    int it;
    int index;
    double *xTmp = malloc(m_order * sizeof(double));
    double *dif = malloc(m_order * sizeof(double));
    double old_b;
    double *old_a;

    old_a = malloc(m_order * sizeof(double));
    old_b = b[row_test];
    for(i = 0; i < m_order; i++)
    {
        old_a[i] = a[m_order*row_test + i];
    }


    for(i = 0; i < m_order_pow_2; i+=(m_order+1))
    {
        temp = i - i%m_order;
        cond = temp + m_order;
        for(j = temp; j < cond; j++)
        {
            if(j != i)
            {
                a[j] /= a[i];
            }
        }
        b[i%m_order] /= a[i];
        a[i] = 0;
    }

    if(row_eval(a,m_order,eps) == 1)
        printf("Convergence test passed\n");
    else {
        printf("Convergence test failed\n");    
        return 0;
    }

    for(it = 0; it < max_it ; it++)
    {
        for(j = 0; j < m_order; j++)
        {
            dif[j] = xTmp[j] = x[j];
        }

        for(j = 0; j < m_order; j++)
        {
            x[j] = b[j];
            index = m_order*j; 
            for(i = 0; i < m_order; i++,index++)
            {
                x[j] -= (a[index]*xTmp[i]);
            }
            dif[j] -= x[j];
        }

        if(max(dif,x,m_order) <= eps)
            break;
    }

    printf("Iterations: %d\n",it);
 
    check_result(row_test,old_a,x,old_b, m_order);

    return x;
}

double max(double *vdif, double *vx, int size)
{
    double max,max2;
    int i;
    double temp,temp2;

    max = max2 = 0;

   for(i = 0; i < size; i++)
   {
        temp = fabs(vdif[i]);
        temp2 = fabs(vx[i]);
        if(temp > max)
            max = temp;
        if(temp2 > max2)
            max2 = temp2;
   }

   return max/max2;
}

int row_eval(double *matrix, int order, double eps)
{
    int i;
    for(i = 0; i < order; i++)
    {
        if(row_convergence(i,matrix,order,eps) == 0)
            return 0;
    }
    return 1;
}

static int row_convergence(int row, double *matrix, int order, double eps)
{
    int row_index = order * row;;
    int i = row_index;;
    double  acc = 0;
    double diff;

    do
    {
        if(i != (row_index + row))
        {
            acc += fabs(matrix[i]);
        }
        i++;
    } while(i < row_index + order);
    diff = 1 - acc;

    return diff > -eps;
}

static void check_result(int row, double *matrix, double *x, double b, int order)
{
    int i = 0;
    double  acc = 0;
    double diff;

    do
    {
        acc += matrix[i]*x[i];   
        i++;
    } while(i < order);
    printf("RowTest: %d => [%lf]=? %lf\n",row,acc,b);
}
